#ifndef SRC_H_MAT_MUL
#define SRC_H_MAT_MUL

#include "cube_signature.hpp"
#include "log_helper.hpp"

#include <helib/helib.h>
#include <omp.h>

#include <array>

#ifdef USE_MPI_WTIME
#include <mpi.h>
// #define omp_get_wtime MPI_Wtime
#endif

// https://portal.nacad.ufrj.br/online/intel/compiler_c/common/core/GUID-7312910C-D175-4544-99C5-29C12D980744.htm
// initializer `omp_priv = omp_orig` using copy constructor of `helib::Ctxt`
#pragma omp declare reduction(ctxt_plus : helib::Ctxt : omp_out += omp_in)     \
    initializer(omp_priv = omp_orig)

#define TOP_TIME_DECL(top)                                                     \
  void start_##top() { time_##top -= omp_get_wtime(); }                        \
  void stop_##top() { time_##top += omp_get_wtime(); }                         \
  double get_time_##top() const { return time_##top; }

struct unit_time {
public:
  TOP_TIME_DECL(unit)

  virtual void reset() { this->time_unit = 0.0; }

protected:
  double time_unit = 0.0;
};

// * `stop_item()` return time increment `delta`
#define TIME_ITEM_DECL(item)                                                   \
  void start_##item() { _start_##item = omp_get_wtime(); }                     \
  double stop_##item() {                                                       \
    double _stop_##item = omp_get_wtime();                                     \
    double delta = _stop_##item - _start_##item;                               \
    time_##item += delta;                                                      \
    return delta;                                                              \
  }                                                                            \
  double get_time_##item() const { return time_##item; }

//! time breakdown in the encoding stage
struct encoding_timer : public unit_time {
public:
  TIME_ITEM_DECL(raw2vec)
  TIME_ITEM_DECL(vec2ptxt)
  TIME_ITEM_DECL(re_ptxt)

  void reset() override {
    unit_time::reset();
    this->time_raw2vec = 0.0;
    this->time_vec2ptxt = 0.0;
    this->time_re_ptxt = 0.0;
  }

protected:
  double time_raw2vec = 0.0;
  double time_vec2ptxt = 0.0;
  double time_re_ptxt = 0.0;

  double _start_raw2vec = 0.0;
  double _start_vec2ptxt = 0.0;
  double _start_re_ptxt = 0.0;
};

//! time breakdown in the homomorphic operation-level
struct hop_timer : public unit_time {
public:
  TIME_ITEM_DECL(rot)
  TIME_ITEM_DECL(add)
  TIME_ITEM_DECL(mult)
  TIME_ITEM_DECL(scale)

  void reset() override {
    unit_time::reset();
    this->time_rot = 0.0;
    this->time_add = 0.0;
    this->time_mult = 0.0;
    this->time_scale = 0.0;
  }

protected:
  double time_rot = 0.0;
  double time_add = 0.0;
  double time_mult = 0.0;
  double time_scale = 0.0;

  double _start_rot = 0.0;
  double _start_add = 0.0;
  double _start_mult = 0.0;
  double _start_scale = 0.0;
};

class hmm_metric_collector;
class constant_handle {
public:
  constant_handle(const helib::EncryptedArray &ea)
      : ea_(ea), threads_(omp_get_max_threads()) {}
  virtual void init(hmm_metric_collector *collector) = 0;

protected:
  const helib::EncryptedArray &ea_;
  const int threads_;
};

class hmm_constant_multipler {
public:
  hmm_constant_multipler() = default;
  virtual ~hmm_constant_multipler() = default;

  virtual void set_no_encoding_reuse(bool val) = 0;

  bool no_encoding_reuse_ = false;
};

enum METHOD { NAIVE, BSGS, METHOD_NUM };
enum MATRIX_TYPE { SQUARE, MAX_K, MIN_K, MEDIUM_K, UNDEF_MATRIX_TYPE };
static constexpr std::array<const char *, UNDEF_MATRIX_TYPE> MATRIX_TYPE_STR = {
    "SQUARE", "MAX_K", "MIN_K", "MEDIUM_K"};

template <typename EngineType> struct hmm_status {
public:
  hmm_status()
      : method_(METHOD::BSGS), threads_(1), reset_(false),
        type_(UNDEF_MATRIX_TYPE) {}
  hmm_status(METHOD method, int threads = 1, bool reset = false)
      : method_(method), threads_(threads), reset_(reset),
        type_(UNDEF_MATRIX_TYPE) {
    ASSERT_PRINTF(threads >= 0, "# of threads must be greater than 0\n");
    ASSERT_PRINTF(method >= 0 && method < METHOD::METHOD_NUM,
                  "Invalid HMM method\n");
  }

  // Getter
  METHOD method() const { return method_; }
  int threads() const { return threads_; }
  bool reset() const { return reset_; }
  MATRIX_TYPE matrix_type() const { return type_; }

  // Setter
  void method(METHOD method) { this->method_ = method; }
  void threads(int threads) { this->threads_ = threads; }
  void reset(bool reset) { this->reset_ = reset; }
  void matrix_type(MATRIX_TYPE type) { this->type_ = type; }

protected:
  METHOD method_; // HMM method: Naive or BSGS
  int threads_;   // # of threads
                  // = 0 uses multi-thread (omp task) version
                  // = 1 uses single-thread version
                  // > 1 uses multi-thread version
  bool reset_;    // whether to reset execution time
  MATRIX_TYPE type_;
};

// * shmm_engine
//   * s: single-ciphertext
//   * hmm: homomorphic matrix multiplication
class shmm_engine {
public:
  static constexpr std::array<const char *, METHOD_NUM> METHOD_STR = {"Naive",
                                                                      "BSGS"};

public:
  static MATRIX_TYPE parse_matrix_type(long m, long n, long k);
  static MATRIX_TYPE parse_matrix_type(const cube_signature &cube);
  static cube_signature get_expected_mnk(long m, long n, long k);

private:
  shmm_engine() = delete;

protected:
  /**
   * @brief Construct a new `shmm_engine` object
   *        * initial setup using (m, n, k)
   *        ! protected:
   *        ! Never be called in the scope out of non-derived class
   * @param m       dim-m of the plaintext matrix
   * @param n       dim-n of the plaintext matrix
   * @param k       dim-k of the plaintext matrix
   * @param status  HMM method(NAIVE/BSGS), # of threads or other status
   */
  shmm_engine(hmm_status<shmm_engine> *status) : status_(status) {}

public:
  /**
   * @brief derived class can override this method to
   *        initialize their own `hmm_constant_multipler`
   *        abd `hmm_metric_collector`
   * @param ea  class for data-movement operations on arrays of slots
   */
  virtual void register_engine(const helib::EncryptedArray &ea,
                               const cube_signature &cube) = 0;

  /**
   * @brief       Execute Ciphertext GEMMs. Single-thread / Multi-thread
   *              Versions are according to the `hmm_status`
   * @param[out]  AB      Enc(A*B)
   * @param[in]   A       Enc(A)
   * @param[in]   B       Enc(B)
   * @param       threads # of threads
   */
  virtual void cgemm(MATRIX_TYPE type, helib::Ctxt &AB, const helib::Ctxt &A,
                     const helib::Ctxt &B) = 0;
  virtual ~shmm_engine();

  virtual hmm_status<shmm_engine> *status() { return status_; }

  virtual void report(std::ostream &os = std::cout);
  virtual void report(MATRIX_TYPE type, std::ostream &os = std::cout);

protected:
  hmm_status<shmm_engine> *status_;
  std::array<hmm_metric_collector *, UNDEF_MATRIX_TYPE> metric_collector_ = {
      nullptr, nullptr, nullptr, nullptr};
  std::array<hmm_constant_multipler *, UNDEF_MATRIX_TYPE> constant_multipler_ =
      {nullptr, nullptr, nullptr, nullptr};

  std::array<const helib::EncryptedArray *, UNDEF_MATRIX_TYPE> ea_;
};

//* unified templates
//! No need to implement
void h_mat_mul(hmm_metric_collector *collector, hmm_constant_multipler *cm,
               const helib::EncryptedArray &ea, helib::Ctxt &AB,
               const helib::Ctxt &A, const helib::Ctxt &B);
void h_mat_mul_mt(hmm_metric_collector *collector, hmm_constant_multipler *cm,
                  int threads, const helib::EncryptedArray &ea, helib::Ctxt &AB,
                  const helib::Ctxt &A, const helib::Ctxt &B);
using hmm_t = decltype(&h_mat_mul);
using hmm_mt_t = decltype(&h_mat_mul_mt);

void hmm_rotate_align(hmm_constant_multipler *hmmsqcm,
                      const helib::EncryptedArray &ea, helib::Ctxt &Mp,
                      const helib::Ctxt &M);
using hmm_ra_t = decltype(&hmm_rotate_align);

class hmm_metric_collector : public unit_time {
public:
  TIME_ITEM_DECL(hmm_encoding)
  TIME_ITEM_DECL(hmm_computing)

  void start_no_encoding_reuse(int tid = 0) {
    if (tid != 0) {
      return;
    }
    no_encoding_reuse -= omp_get_wtime();
  }
  void stop_no_encoding_reuse(int tid = 0) {
    if (tid != 0) {
      return;
    }
    no_encoding_reuse += omp_get_wtime();
  }

  double get_no_encoding_reuse() const { return no_encoding_reuse; }

  void reset() override {
    unit_time::reset();
    reset_encoding();
    reset_computing();
  }

  void reset_encoding() {
    this->time_hmm_encoding = 0.0;
    this->time_raw2vec = 0.0;
    this->time_vec2ptxt = 0.0;
    this->time_re_ptxt = 0.0;
  }

  virtual void reset_computing() {
    //! clear `time_hmm_computing` from `time_unit`
    this->time_unit -= this->time_hmm_computing;

    // clear to 0
    this->no_encoding_reuse = 0.0;
    this->time_hmm_computing = 0.0;
    this->time_rot = 0.0;
    this->time_add = 0.0;
    this->time_mult = 0.0;
    this->time_scale = 0.0;
  }

  virtual void report(const hmm_status<shmm_engine> *status,
                      std::ostream &os = std::cout) = 0;
  virtual ~hmm_metric_collector() = default;

protected:
  // * Overall breakdown
  //   * time_hmm_encoding (Encoding Stage)
  //     * time_raw2vec    (raw data -> vector)
  //     * time_vec2ptxt   (vector -> plaintext)
  //     * time_re_ptxt    (plaintext -> plaintext)
  double time_hmm_encoding = 0.0;
  double time_raw2vec = 0.0;
  double time_vec2ptxt = 0.0;
  double time_re_ptxt = 0.0;
  //   * time_hmm_computing (Computing Stage)
  //     * time_rot         (Rotation)
  //     * time_add         (Ctxt-Ctxt Addition)
  //     * time_mult        (Ctxt-Ctxt Multiplication)
  //     * time_scale       (Ptxt-Ctxt Multiplication)
  double time_hmm_computing = 0.0;
  double time_rot = 0.0;
  double time_add = 0.0;
  double time_mult = 0.0;
  double time_scale = 0.0;

  double no_encoding_reuse = 0.0;

  double _start_hmm_encoding = 0.0;
  double _start_hmm_computing = 0.0;
};

//* Some util marcos for class `hmm_metric_collector`
//* 1. `subprogress` is derived from `unit_time`
//* 2. `subprogress` has its own time entry `time_##item`
//* 3. `subprogress` only call `start_##subprogress()` and
//*    `stop_##subprogress()` once
//*
//* these marcos
//* 1. use each `subprogress`'s `time_##item` to update
//*    overall `time_##item` in `hmm_metric_collector`
//* 2. use each `subprogress`'s `time_unit` to update
//*    overall `time_unit` in `hmm_metric_collector`
//? 3. There may be problems with multi-threading
//*
#define HMM_UNIT_TIME_UTIL_DECL(subprogress)                                   \
  void start_##subprogress() { subprogress.start_unit(); }                     \
  void stop_##subprogress() {                                                  \
    subprogress.stop_unit();                                                   \
    time_unit += subprogress.get_time_unit();                                  \
  }
#define HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, item)                        \
  void start_##subprogress##_##item(int tid = 0) {                             \
    if (tid != 0) {                                                            \
      return;                                                                  \
    }                                                                          \
    subprogress.start_##item();                                                \
  }                                                                            \
  void stop_##subprogress##_##item(int tid = 0) {                              \
    if (tid != 0) {                                                            \
      return;                                                                  \
    }                                                                          \
    time_##item += subprogress.stop_##item();                                  \
  }
#define ENCODING_TIME_UTIL_DECL(subprogress)                                   \
  HMM_UNIT_TIME_UTIL_DECL(subprogress)                                         \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, raw2vec)                           \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, vec2ptxt)                          \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, re_ptxt)
#define HOP_TIME_UTIL_DECL(subprogress)                                        \
  HMM_UNIT_TIME_UTIL_DECL(subprogress)                                         \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, rot)                               \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, add)                               \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, mult)                              \
  HMM_UNIT_ITEM_TIME_UTIL_DECL(subprogress, scale)

hmm_metric_collector *create_metric(MATRIX_TYPE type);

// For homomorphic matrix multiplication
class hmmsq_collector : public hmm_metric_collector {
public:
  HOP_TIME_UTIL_DECL(computing_ra)
  HOP_TIME_UTIL_DECL(computing_sc)

  ENCODING_TIME_UTIL_DECL(encoding_ra)
  ENCODING_TIME_UTIL_DECL(encoding_sc)

  void reset() override {
    hmm_metric_collector::reset();
    encoding_ra.reset();
    encoding_sc.reset();
    computing_ra.reset();
    computing_sc.reset();
  }
  void reset_computing() override {
    computing_ra.reset();
    computing_sc.reset();
    hmm_metric_collector::reset_computing();
  }

  void report(const hmm_status<shmm_engine> *status,
              std::ostream &os = std::cout) override;

private:
  // * Stage-by-stage breakdown
  //   * Encoding Stage
  encoding_timer encoding_ra;
  encoding_timer encoding_sc;
  //   * Computing Stage
  hop_timer computing_ra;
  hop_timer computing_sc;

public:
  // noise
  double left_cap = 0.0;
  double final_noise = 0.0;
  double init_noise = 0.0;
  double ra_noise = 0.0;
  double ra_correct_noise = 0.0;
  double sc_noise = 0.0;

  // err
  double final_err = 0.0;
  double init_err = 0.0;
  double ra_err = 0.0;
  double sc_rot_err = 0.0;
  double sc_hmul_err = 0.0;
};

class hmm_maxk_collector : public hmm_metric_collector {
public:
  HOP_TIME_UTIL_DECL(computing_rr)
  HOP_TIME_UTIL_DECL(computing_ra)
  HOP_TIME_UTIL_DECL(computing_sc)
  HOP_TIME_UTIL_DECL(computing_sa)

  ENCODING_TIME_UTIL_DECL(encoding_ra)
  ENCODING_TIME_UTIL_DECL(encoding_sc)

  void reset() override {
    hmm_metric_collector::reset();
    encoding_ra.reset();
    encoding_sc.reset();
    computing_rr.reset();
    computing_ra.reset();
    computing_sc.reset();
    computing_sa.reset();
  }
  void reset_computing() override {
    computing_rr.reset();
    computing_ra.reset();
    computing_sc.reset();
    computing_sa.reset();
    hmm_metric_collector::reset_computing();
  }

  void report(const hmm_status<shmm_engine> *status,
              std::ostream &os = std::cout) override;

private:
  // * Stage-by-stage breakdown
  //   * Encoding Stage
  encoding_timer encoding_ra;
  encoding_timer encoding_sc;
  //   * Computing Stage
  hop_timer computing_rr;
  hop_timer computing_ra;
  hop_timer computing_sc;
  hop_timer computing_sa;

public:
  // noise
  double left_cap = 0.0;
  double final_noise = 0.0;
  double init_noise = 0.0;
  double rr_noise = 0.0;
  double ra_noise = 0.0;
  double ra_correct_noise = 0.0;
  double sc_noise = 0.0;
  double sa_noise = 0.0;

  // err
  double final_err = 0.0;
  double init_err = 0.0;
  double rr_err = 0.0;
  double ra_err = 0.0;
  double sc_rot_err = 0.0;
  double sc_hmul_err = 0.0;
  double sa_err = 0.0;
};

class hmm_mink_collector : public hmm_metric_collector {
public:
  HOP_TIME_UTIL_DECL(computing_rr)
  HOP_TIME_UTIL_DECL(computing_ra)
  HOP_TIME_UTIL_DECL(computing_sc)

  ENCODING_TIME_UTIL_DECL(encoding_ra)
  ENCODING_TIME_UTIL_DECL(encoding_sc)

  void reset() override {
    hmm_metric_collector::reset();
    encoding_ra.reset();
    encoding_sc.reset();
    computing_rr.reset();
    computing_ra.reset();
    computing_sc.reset();
  }
  void reset_computing() override {
    computing_rr.reset();
    computing_ra.reset();
    computing_sc.reset();
    hmm_metric_collector::reset_computing();
  }

  void report(const hmm_status<shmm_engine> *status,
              std::ostream &os = std::cout) override;

private:
  // * Stage-by-stage breakdown
  //   * Encoding Stage
  encoding_timer encoding_ra;
  encoding_timer encoding_sc;
  //   * Computing Stage
  hop_timer computing_rr;
  hop_timer computing_ra;
  hop_timer computing_sc;

public:
  // noise
  double left_cap = 0.0;
  double final_noise = 0.0;
  double init_noise = 0.0;
  double rr_noise = 0.0;
  double ra_noise = 0.0;
  double ra_correct_noise = 0.0;
  double sc_noise = 0.0;

  // err
  double final_err = 0.0;
  double init_err = 0.0;
  double rr_err = 0.0;
  double ra_err = 0.0;
  double sc_rot_err = 0.0;
  double sc_hmul_err = 0.0;
};

#endif /* SRC_HYPERCUBE_H_MAT_MUL */
